import warnings

from omegaconf import DictConfig

import neurokit2 as nk
import numpy as np

# Suppress warnings
warnings.filterwarnings('ignore')
from sklearn.cluster import DBSCAN
from scipy.stats import mode


class FindEDES:
    """
    This class finds ED and ES points from a 12-lead ECG.
    """
    def __init__(self, cfg: DictConfig, input_electrodes: int = 12, sampling_rate: int = 500):
        self.sampling_rate = sampling_rate
        self.input_electrodes = input_electrodes

        self.denoised = cfg['clean_ecg']['enable']
        self.clean_method = cfg['clean_ecg']['method']
        self.delta_r = cfg['es_peak_find']['delta_r']
        self.delta_t = cfg['es_peak_find']['delta_t']
        self.peak_find_method = cfg['es_peak_find']['method']

        self.min_peaks_in_cluster_r = cfg['es_peak_find']['min_peaks_in_cluster_r']
        self.min_peaks_in_cluster_t = cfg['es_peak_find']['min_peaks_in_cluster_t']

        self.variance_threshold = 0.01
        self.segment_length = 500

    def _check_variance_consistency(self, ecg):
        overall_variance = np.var(ecg)
        segment_variances = [
            np.var(ecg[i:i+self.segment_length])
            for i in range(0, len(ecg), self.segment_length)
        ]
        variance_threshold = overall_variance * 0.1
        low_variance_segments = sum(var < variance_threshold for var in segment_variances)
        total_segments = len(segment_variances)
        return low_variance_segments / total_segments < 0.5 

    def _find_representative_peaks(self, peaks_all_channels, delta, min_samples):
        all_peaks = np.sort(np.concatenate(peaks_all_channels))
        if all_peaks.size == 0:
            return np.array([])

        all_peaks_reshaped = all_peaks.reshape(-1, 1)
        db = DBSCAN(eps=delta, min_samples=min_samples).fit(all_peaks_reshaped)
        labels = db.labels_

        unique_labels = np.unique(labels[labels >= 0])
        representatives = []
        for label in unique_labels:
            cluster_peaks = all_peaks[labels == label]
            peak_mode = mode(cluster_peaks)
            if peak_mode.count.size > 0 and peak_mode.count >= min_samples:
                most_frequent_peak = peak_mode.mode
            else:
                most_frequent_peak = np.median(cluster_peaks)
            representatives.append(int(most_frequent_peak))

        return np.array(representatives)

    def _calculate_mean_ratio(self, aggregated_rpeaks, aggregated_tpeaks):
        r_to_t_intervals = []
        for rpeak in aggregated_rpeaks:
            following_tpeaks = aggregated_tpeaks[aggregated_tpeaks > rpeak]
            if following_tpeaks.size > 0:
                r_to_t_intervals.append(following_tpeaks[0] - rpeak)
        
        r_to_r_intervals = np.diff(aggregated_rpeaks)
        
        min_length = min(len(r_to_t_intervals), len(r_to_r_intervals))
        r_to_t_intervals = np.array(r_to_t_intervals[:min_length])
        r_to_r_intervals = r_to_r_intervals[:min_length]

        ratios = np.where(r_to_r_intervals != 0,
                          r_to_t_intervals / r_to_r_intervals,
                          np.nan)

        return np.nanmean(ratios)

    def _get_rpeaks(self, ecg):
        try:
            peaks_info = nk.ecg_peaks(ecg, sampling_rate=self.sampling_rate, correct_artifacts=True, method=self.peak_find_method)[1]
            peaks = peaks_info['ECG_R_Peaks']
        except Exception:
            peaks = np.array([])

        if len(peaks) < 3: 
            try:
                peaks_info = nk.ecg_peaks(ecg, sampling_rate=self.sampling_rate, correct_artifacts=True, method="koka2022")[1]
                peaks_fallback = peaks_info['ECG_R_Peaks']
                return peaks_fallback
            except Exception as e:
                return np.array([])
        return peaks

    def _get_tpeaks(self, ecg, rpeaks):
        try:
            tpeaks_info = nk.ecg_delineate(ecg, rpeaks, check=False, sampling_rate=self.sampling_rate, method="cwt")[1]
            tpeaks = np.array(tpeaks_info['ECG_T_Peaks'], dtype=float)
            tpeaks = tpeaks[~np.isnan(tpeaks)]
            tpeaks = tpeaks.astype(int)
        except Exception:
            tpeaks = np.array([])

        if len(tpeaks) < 3:  # Use fallback method if fewer than 3 peaks detected
            try:
                tpeaks_info_fallback = nk.ecg_delineate(ecg, rpeaks, sampling_rate=self.sampling_rate, method="peak")[1]
                tpeaks_fallback = np.array(tpeaks_info_fallback['ECG_T_Peaks'], dtype=float)
                tpeaks_fallback = tpeaks_fallback[~np.isnan(tpeaks_fallback)]
                return tpeaks_fallback.astype(int)
            except Exception:
                return np.array([])
        return tpeaks

    def __call__(self, sample: np.ndarray) -> dict:
        rpeaks_all_channels = []
        tpeaks_all_channels = []

        for i, ecg in enumerate(sample[:3]):

            ecg = ecg if self.denoised else nk.ecg_clean(ecg, sampling_rate=self.sampling_rate, method=self.clean_method)
            if not self._check_variance_consistency(ecg):
                continue 

            rpeaks = self._get_rpeaks(ecg)
            rpeaks = rpeaks[~np.isnan(rpeaks)]

            print(f"RPEAKS LEAD {i} - {rpeaks}")
            if rpeaks.size == 0:
                continue
            rpeaks_all_channels.append(rpeaks)

            tpeaks = self._get_tpeaks(ecg, rpeaks)
            
            print(f"TPEAKS LEAD {i} - {tpeaks}")

            if tpeaks.size > 0:
                tpeaks_all_channels.append(tpeaks)

        aggregated_rpeaks = self._find_representative_peaks(rpeaks_all_channels,
                                                            delta=self.delta_r, min_samples=self.min_peaks_in_cluster_r)
        aggregated_tpeaks = self._find_representative_peaks(tpeaks_all_channels,
                                                            delta=self.delta_t, min_samples=self.min_peaks_in_cluster_t)

        aggregated_rpeaks = np.sort(aggregated_rpeaks)
        aggregated_tpeaks = np.sort(aggregated_tpeaks)

        mean_ratio = self._calculate_mean_ratio(aggregated_rpeaks, aggregated_tpeaks)

        return {
            "ecg": sample,
            "rpeaks": np.atleast_1d(aggregated_rpeaks),
            "tpeaks": np.atleast_1d(aggregated_tpeaks),
            "ratio": np.atleast_1d(mean_ratio)
        }